{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬▬]                                   Elapsed time: 5s (2.0 it/s)            \n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "bc7ab4a5b7f548a081c7952e6e324af5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Plot(antialias=3, axes=['x', 'y', 'z'], background_color=16777215, camera=[4.5, 4.5, 4.5, 0.0, 0.0, 0.0, 1.0, …"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "\"\"\"Solution of a particular nonlinear time-dependent \n",
    "fourth-order equation, known as the Cahn-Hilliard equation.\"\"\"\n",
    "import random\n",
    "from dolfin import *\n",
    "set_log_level(30)\n",
    "\n",
    "# Class representing the intial conditions\n",
    "class InitialConditions(UserExpression):\n",
    "    def __init__(self, **kwargs):\n",
    "        random.seed(2 + MPI.rank(MPI.comm_world))\n",
    "        super().__init__(**kwargs)\n",
    "    def eval(self, values, x):\n",
    "        values[0] = 0.63 + 0.02 * (0.5 - random.random())\n",
    "        values[1] = 0.0\n",
    "    def value_shape(self): return (2,)\n",
    "\n",
    "# Class for interfacing with the Newton solver\n",
    "class CahnHilliardEquation(NonlinearProblem):\n",
    "    def __init__(self, a, L):\n",
    "        NonlinearProblem.__init__(self)\n",
    "        self.L = L\n",
    "        self.a = a\n",
    "    def F(self, b, x): assemble(self.L, tensor=b)\n",
    "    def J(self, A, x): assemble(self.a, tensor=A)\n",
    "\n",
    "# Model parameters\n",
    "lmbda = 1.0e-02  # surface parameter\n",
    "dt    = 5.0e-06  # time step\n",
    "# time stepping family, \n",
    "# e.g. theta=1 -> backward Euler, theta=0.5 -> Crank-Nicolson\n",
    "theta = 0.5\n",
    "\n",
    "# Form compiler options\n",
    "parameters[\"form_compiler\"][\"optimize\"] = True\n",
    "parameters[\"form_compiler\"][\"cpp_optimize\"] = True\n",
    "\n",
    "# Create mesh and define function spaces\n",
    "mesh = UnitSquareMesh(60, 60)\n",
    "# mesh = UnitSquareMesh.create(60, 60, CellType.Type.triangle)\n",
    "# V = FunctionSpace(mesh, \"Lagrange\", 1)\n",
    "P1 = FiniteElement(\"Lagrange\", mesh.ufl_cell(), 1)\n",
    "ME = FunctionSpace(mesh, P1 * P1)\n",
    "\n",
    "# Define trial and test functions\n",
    "du   = TrialFunction(ME)\n",
    "q, v = TestFunctions(ME)\n",
    "\n",
    "# Define functions\n",
    "u  = Function(ME)  # current solution\n",
    "u0 = Function(ME)  # solution from previous converged step\n",
    "\n",
    "# Split mixed functions\n",
    "dc, dmu = split(du)\n",
    "c, mu   = split(u)\n",
    "c0, mu0 = split(u0)\n",
    "\n",
    "# Create intial conditions and interpolate\n",
    "u_init = InitialConditions(degree=1)\n",
    "u.interpolate(u_init)\n",
    "u0.interpolate(u_init)\n",
    "\n",
    "# Compute the chemical potential df/dc\n",
    "c = variable(c)\n",
    "f = 100 * c ** 2 * (1 - c) ** 2\n",
    "mu_mid = (1 - theta) * mu0 + theta * mu\n",
    "\n",
    "# Weak statement of the equations\n",
    "L0 = c * q - c0 * q + dt * dot(grad(mu_mid), grad(q))\n",
    "L1 = mu * v - diff(f, c) * v - lmbda * dot(grad(c), grad(v))\n",
    "L  = (L0 + L1) * dx\n",
    "\n",
    "# Compute directional derivative about u in the direction of du\n",
    "a = derivative(L, u, du)\n",
    "\n",
    "# Create nonlinear problem and Newton solver\n",
    "problem = CahnHilliardEquation(a, L)\n",
    "solver = NewtonSolver()\n",
    "solver.parameters[\"linear_solver\"] = \"lu\"\n",
    "solver.parameters[\"convergence_criterion\"] = \"incremental\"\n",
    "solver.parameters[\"relative_tolerance\"] = 1e-6\n",
    "\n",
    "# Step in time\n",
    "from vtkplotter.dolfin import ProgressBar, plot\n",
    "\n",
    "t = 0\n",
    "pb = ProgressBar(0,10)\n",
    "for i in pb.range():\n",
    "    t += dt\n",
    "    u0.vector()[:] = u.vector()\n",
    "    solver.solve(problem, u.vector())\n",
    "    \n",
    "    plt = plot(u.split()[0], z=i*0.1, cmap='hot', add=True) # do not clear the canvas\n",
    "    pb.print()\n",
    "         \n",
    "plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
